import pandas as pd
from sqlalchemy import create_engine, Table, Column, Integer, String, Text, MetaData, ForeignKey

# 连接到数据库
DATABASE_URI = 'mysql+pymysql://root:123456@localhost:3306/ai_risk'
engine = create_engine(DATABASE_URI)
metadata = MetaData()

# 定义表结构
# （表结构定义代码保持不变）
risks = Table('risks', metadata,
    Column('id', Integer, primary_key=True),
    Column('title', String(255), nullable=False),
    Column('ev_id', String(50), unique=True),
    Column('description', Text),
    Column('additional_evidence', Text),
    Column('causal_factor_id', Integer, ForeignKey('causal_factors.id')),
    Column('domain_id', Integer, ForeignKey('domains.id')),
    Column('sub_domain_id', Integer, ForeignKey('sub_domains.id')),
    Column('category_id', Integer, ForeignKey('risk_categories.id')),
    Column('sub_category_id', Integer, ForeignKey('sub_categories.id'))
)

causal_factors = Table('causal_factors', metadata,
    Column('id', Integer, primary_key=True),
    Column('entity_type', String(50)),
    Column('intent_type', String(50)),
    Column('timing_type', String(50))
)

domains = Table('domains', metadata,
    Column('id', Integer, primary_key=True),
    Column('domain_name', String(255), nullable=False)
)

sub_domains = Table('sub_domains', metadata,
    Column('id', Integer, primary_key=True),
    Column('domain_id', Integer, ForeignKey('domains.id')),
    Column('sub_domain_name', String(255), nullable=False)
)

risk_categories = Table('risk_categories', metadata,
    Column('id', Integer, primary_key=True),
    Column('category_level', String(50)),
    Column('category_name', String(255))
)

sub_categories = Table('sub_categories', metadata,
    Column('id', Integer, primary_key=True),
    Column('category_id', Integer, ForeignKey('risk_categories.id')),
    Column('sub_category_name', String(255))
)

# 创建表
metadata.create_all(engine)

# 读取Excel文件
file_path = './Ai risk database.xlsx'
df = pd.read_excel(file_path, sheet_name='Sheet1')

# 用 None 替换 DataFrame 中的 NaN 值
df = df.where(pd.notnull(df), None)

# 数据清洗和插入
with engine.connect() as conn:
    trans = conn.begin()  # 开始事务
    try:
        for _, row in df.iterrows():
            # 插入或获取因果分类ID
            causal_factor = conn.execute(
                causal_factors.select().where(
                    causal_factors.c.entity_type == row['Entity'],
                    causal_factors.c.intent_type == row['Intent'],
                    causal_factors.c.timing_type == row['Timing']
                )
            ).mappings().fetchone()

            if causal_factor is None:
                causal_factor_id = conn.execute(causal_factors.insert().values(
                    entity_type=row['Entity'],
                    intent_type=row['Intent'],
                    timing_type=row['Timing']
                )).inserted_primary_key[0]
            else:
                causal_factor_id = causal_factor['id']

            # 插入或获取领域ID
            domain = conn.execute(domains.select().where(domains.c.domain_name == row['Domain'])).mappings().fetchone()
            if domain is None:
                domain_id = conn.execute(domains.insert().values(domain_name=row['Domain'])).inserted_primary_key[0]
            else:
                domain_id = domain['id']

            # 插入或获取子领域ID
            sub_domain = conn.execute(sub_domains.select().where(
                sub_domains.c.domain_id == domain_id,
                sub_domains.c.sub_domain_name == row['Sub-domain']
            )).mappings().fetchone()
            if sub_domain is None:
                sub_domain_id = conn.execute(sub_domains.insert().values(
                    domain_id=domain_id,
                    sub_domain_name=row['Sub-domain']
                )).inserted_primary_key[0]
            else:
                sub_domain_id = sub_domain['id']

            # 插入或获取风险类别ID
            category = conn.execute(risk_categories.select().where(
                risk_categories.c.category_level == row['Category level'],
                risk_categories.c.category_name == row['Risk category']
            )).mappings().fetchone()
            if category is None:
                category_id = conn.execute(risk_categories.insert().values(
                    category_level=row['Category level'],
                    category_name=row['Risk category']
                )).inserted_primary_key[0]
            else:
                category_id = category['id']

            # 插入或获取子类别ID
            sub_category = conn.execute(sub_categories.select().where(
                sub_categories.c.category_id == category_id,
                sub_categories.c.sub_category_name == row['Risk subcategory']
            )).mappings().fetchone()
            if sub_category is None:
                sub_category_id = conn.execute(sub_categories.insert().values(
                    category_id=category_id,
                    sub_category_name=row['Risk subcategory']
                )).inserted_primary_key[0]
            else:
                sub_category_id = sub_category['id']

            # 检查是否已存在该 ev_id
            existing_risk = conn.execute(risks.select().where(risks.c.ev_id == row['Ev_ID'])).mappings().fetchone()
            if existing_risk is None:
                # 如果不存在，则插入风险信息
                conn.execute(risks.insert().values(
                    title=row['Title'],
                    ev_id=row['Ev_ID'],
                    description=row['Description'],
                    additional_evidence=row['Additional ev.'],
                    causal_factor_id=causal_factor_id,
                    domain_id=domain_id,
                    sub_domain_id=sub_domain_id,
                    category_id=category_id,
                    sub_category_id=sub_category_id
                ))
            else:
                # 更新已有记录
                conn.execute(risks.update().where(risks.c.ev_id == row['Ev_ID']).values(
                    title=row['Title'],
                    description=row['Description'],
                    additional_evidence=row['Additional ev.'],
                    causal_factor_id=causal_factor_id,
                    domain_id=domain_id,
                    sub_domain_id=sub_domain_id,
                    category_id=category_id,
                    sub_category_id=sub_category_id
                ))
                print(f"Updated entry for ev_id {row['Ev_ID']}.")

        trans.commit()  # 提交事务
    except Exception as e:
        trans.rollback()  # 回滚事务
        print(f"An error occurred: {e}")



